import uuid
from typing import Dict

from sqlalchemy import Column, Integer, String, Boolean, SmallInteger, ForeignKey, Text, JSON

from db.base import Base
from db.models.base import ARRAY, GUID, BaseDatabaseModel


class Conversation(Base, BaseDatabaseModel):
	__tablename__ = "conversation"

	id = Column(GUID(), primary_key=True, index=True, default=uuid.uuid1)
	question = Column(Text, nullable=True, comment="quesion")
	answer = Column(Text, nullable=True, comment="answer")
	image_path = Column(String(512), nullable=True, default=None, comment="图片问答时的图片名")
	reference = Column(ARRAY(), nullable=True, default=[])  # 存储可以访问的用户组ID列表，以逗号分隔
	is_delete = Column(Boolean, index=True, default=False)
	file_id = Column(String(512), nullable=True, default='', comment="对于应急和双控来讲，存在并且是文件id")

	dialog_id = Column(GUID(), ForeignKey('dialog.id'))

	def to_json(self):
		return {
			"id": str(self.id),
			"question": self.question,
			"answer": self.answer,
			"image": self.image_path,
			"reference": self.reference,
			"updated_time": self.updated_time.strftime("%Y-%m-%d %H:%M:%S"),
		}

	def to_messages(self):
		import re
		answer = re.sub(r'<cite>[^<]*<\/cite>', '', str(self.answer))
		history = [{"role": "user", "content": self.question}, {"role": "assistant", "content": answer}]
		return history


class Dialog(BaseDatabaseModel, Base):
	__tablename__ = "dialog"

	id = Column(GUID(), primary_key=True, index=True, default=uuid.uuid1)
	name = Column(String(255), nullable=True, index=True, comment="dialog application name")
	description = Column(String(2048), nullable=True,  comment="dialog description")

	kb_ids = Column(ARRAY(), nullable=False, default=[])
	is_delete = Column(Boolean, index=True, default=False)
	icon = Column(Text, nullable=True, comment="icon base64 string")
	llm_id = Column(String(128), nullable=True, comment="default llm ID")
	dialog_type = Column(SmallInteger, nullable=True, index=True, comment="对话类型，0:原生模型对话,1:知识问答/2:班前会/3:应急预案/4:事故案例")

	user_id = Column(Integer, ForeignKey('users.id'))

	def to_json(self):
		return {
			"name": self.name,
			"id": str(self.id),
			"dialog_type": self.dialog_type,
			"updated_time": self.updated_time.strftime("%Y-%d-%m %H:%M:%S"),
			"kb_ids": self.kb_ids
		}